# 以下库需要在镜像里安装好
import config as configs
import os
import shutil
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

print("TensorFlow version:", tf.__version__)

parser = configs.get_parser()
args = parser.parse_args()

print(args)

curpath = os.getcwd()
print(curpath)
modelpath = None
datapath = None

# 使用界面选取的模型可参考下面的示例
# if os.path.isfile(args.model_load_dir):
#     modelpath = args.model_load_dir
# else:
#     for file in os.listdir(args.model_load_dir):
#         modelpath = args.model_load_dir + '/' + file
        
# model = YOLO(curpath + "/yolov8n.pt")

# mnist = tf.keras.datasets.mnist
# (x_train, y_train), (x_test, y_test) = mnist.load_data()
if os.path.isfile(args.data_url):
    datapath = args.data_url
else:
    for file in os.listdir(args.data_url):
        datapath = args.data_url + '/' + file

with np.load(datapath) as datas:
    x_train, y_train = datas['x_train'], datas['y_train']
    x_test, y_test = datas['x_test'], datas['y_test']
    
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

predictions = model(x_train[:1]).numpy()

tf.nn.softmax(predictions).numpy()

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

loss_fn(y_train[:1], predictions).numpy()

model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)

model.evaluate(x_test,  y_test, verbose=2)

probability_model = tf.keras.Sequential([
  model,
  tf.keras.layers.Softmax()
])

probability_model(x_test[:5])

model.export(args.train_model_out + '/my_model')